#!/usr/bin/env python3
"""
Claude Code User Setup Script
Automatically configure user authentication and Keychain credentials based on email
"""

import subprocess
import json
import sys
import re
import secrets
import requests
from datetime import datetime
import urllib3
import os
import platform

# Disable SSL warnings (if using local proxy)
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# Configuration
PROXY_BASE_URL = "https://claudeproxy.corp.astratech.ae"  # Proxy server address
VERIFY_SSL = False  # Whether to verify SSL certificate
PROXY_HOST = "claudeproxy.corp.astratech.ae"
PROXY_IP = "10.44.14.254"

def validate_email(email):
    """Validate email format"""
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
    return re.match(pattern, email) is not None


def check_connectivity():
    """Check connectivity to proxy server"""
    print("Step 2: Checking network connectivity...")
    
    # Use subprocess to call curl command for checking
    try:
        # Use curl to check connectivity
        test_url = f"https://{PROXY_HOST}/"
        curl_cmd = [
            'curl',
            '-s',                    # Silent mode
            '--connect-timeout', '5', # Connection timeout 5 seconds
            '--max-time', '10',       # Max execution time 10 seconds
            '-o', '/dev/null',        # Discard output
            '-w', '%{http_code}',     # Only output HTTP status code
            '-k',                     # Ignore SSL certificate verification
            test_url
        ]
        
        result = subprocess.run(curl_cmd, capture_output=True, text=True)
        http_code = result.stdout.strip()
        
        # Any non-000 HTTP status code indicates successful connection
        if http_code and http_code != '000':
            print(f"  ✅ Successfully connected to {PROXY_HOST} (HTTP {http_code})")
            return True
        else:
            # If first attempt fails, try different endpoints
            test_urls = [
                f"https://{PROXY_HOST}/health",
                f"https://{PROXY_HOST}/api/users"
            ]
            
            for url in test_urls:
                curl_cmd[-1] = url  # Update URL
                result = subprocess.run(curl_cmd, capture_output=True, text=True)
                http_code = result.stdout.strip()
                
                if http_code and http_code != '000':
                    print(f"  ✅ Successfully connected to {PROXY_HOST} (HTTP {http_code})")
                    return True
            
            # All attempts failed
            print(f"  ❌ Unable to connect to {PROXY_HOST}")
            print("  ⚠️ Please check:")
            print("     1. Company VPN is connected")
            print("     2. Proxy service is running")
            print("     3. Firewall settings are not blocking connection")
            print()
            print("  Tip: Please connect to company VPN and rerun this script")
            return False
            
    except FileNotFoundError:
        print("  ❌ curl command not found, please install curl first")
        return False
    except Exception as e:
        print(f"  ❌ Connection check failed: {e}")
        return False

def get_keychain_credentials():
    """Read Claude Code credentials from Keychain"""
    try:
        result = subprocess.run([
            'security', 
            'find-generic-password',
            '-s', 'Claude Code-credentials',
            '-w'
        ], capture_output=True, text=True, check=True)
        
        password_data = result.stdout.strip()
        
        try:
            credentials = json.loads(password_data)
            return credentials
        except json.JSONDecodeError:
            return password_data
            
    except subprocess.CalledProcessError:
        return None

def write_to_keychain(credentials_json):
    """Write credentials to Keychain"""
    try:
        # Get current system logged-in username
        import os
        username = os.getenv('USER') or os.getenv('USERNAME') or 'Claude Code'
        
        # First try to delete existing password item (if exists)
        subprocess.run([
            'security',
            'delete-generic-password',
            '-s', 'Claude Code-credentials'
        ], capture_output=True, check=False)  # Ignore deletion failure errors
        
        # Add new password item using current username
        # -U parameter allows all applications to access this password item
        result = subprocess.run([
            'security',
            'add-generic-password',
            '-s', 'Claude Code-credentials',
            '-a', username,  # Use current system username
            '-w', credentials_json,  # Password content
            '-U'  # Allow all applications to access (Update mode - no access control)
        ], capture_output=True, text=True, check=True)
        
        return True
    except subprocess.CalledProcessError as e:
        print(f"Failed to write to Keychain: {e.stderr}")
        return False

def add_oauth_token(name, token_json):
    """Add OAuth token to proxy server"""
    try:
        url = f"{PROXY_BASE_URL}/api/oauth/tokens"
        payload = {
            "name": name,
            "token": token_json if isinstance(token_json, str) else json.dumps(token_json),
            "force": True  # Force overwrite duplicate tokens
        }
        
        response = requests.post(url, json=payload, verify=VERIFY_SSL)
        
        if response.status_code == 200 or response.status_code == 409:
            result = response.json()
            if response.status_code == 409:
                print(f"  Duplicate token detected, using force mode to overwrite")
                # Retry with force parameter
                payload["force"] = True
                response = requests.post(url, json=payload, verify=VERIFY_SSL)
                result = response.json()
            
            print(f"  ✅ OAuth token added successfully: {result.get('name', name)}")
            return True
        else:
            print(f"  ❌ Failed to add OAuth token: {response.status_code}")
            print(f"     {response.text}")
            return False
    except Exception as e:
        print(f"  ❌ Failed to call OAuth API: {e}")
        return False

def get_user_list():
    """Get user list"""
    try:
        url = f"{PROXY_BASE_URL}/api/users"
        response = requests.get(url, verify=VERIFY_SSL)
        
        if response.status_code == 200:
            return response.json()
        else:
            print(f"  ⚠️ Failed to get user list: {response.status_code}")
            return []
    except Exception as e:
        print(f"  ⚠️ Failed to call user API: {e}")
        return []

def add_user(name, token):
    """Add user to user list"""
    try:
        url = f"{PROXY_BASE_URL}/api/users"
        payload = {
            "name": name,
            "token": token
        }
        
        response = requests.post(url, json=payload, verify=VERIFY_SSL)
        
        if response.status_code == 200:
            result = response.json()
            print(f"  ✅ User added successfully: {result.get('name', name)}")
            return result
        else:
            print(f"  ❌ Failed to add user: {response.status_code}")
            print(f"     {response.text}")
            return None
    except Exception as e:
        print(f"  ❌ Failed to call user API: {e}")
        return None

def generate_user_token():
    """Generate 60-digit hex string as user token"""
    return secrets.token_hex(30)  # 30 bytes = 60 hex characters

def create_keychain_credentials(user_token):
    """Create Keychain credentials JSON"""
    credentials = {
        "claudeAiOauth": {
            "accessToken": user_token,
            "refreshToken": user_token,
            "expiresAt": 1856184395,  # Fixed future timestamp
            "scopes": [
                "user:inference",
                "user:profile"
            ],
            "subscriptionType": "max"
        }
    }
    return json.dumps(credentials)

def check_and_set_env_variable():
    """Check and set ANTHROPIC_BASE_URL environment variable"""
    print("Step 6: Checking environment variable configuration...")
    
    current_value = os.environ.get('ANTHROPIC_BASE_URL')
    
    if current_value == PROXY_BASE_URL:
        print(f"  ✅ ANTHROPIC_BASE_URL is correctly set to: {PROXY_BASE_URL}")
        return True
    
    if current_value:
        print(f"  ⚠️ Current ANTHROPIC_BASE_URL: {current_value}")
        print(f"  Need to update to: {PROXY_BASE_URL}")
    else:
        print(f"  ℹ️ ANTHROPIC_BASE_URL not set")
        print(f"  Need to set to: {PROXY_BASE_URL}")
    
    # Determine shell configuration file to update
    shell = os.environ.get('SHELL', '/bin/bash')
    home = os.path.expanduser('~')
    
    # Choose configuration file based on shell
    if 'zsh' in shell:
        config_files = [f'{home}/.zshrc', f'{home}/.zprofile']
    elif 'bash' in shell:
        config_files = [f'{home}/.bashrc', f'{home}/.bash_profile', f'{home}/.profile']
    else:
        config_files = [f'{home}/.profile']
    
    # Environment variable line to add
    export_line = f'export ANTHROPIC_BASE_URL="{PROXY_BASE_URL}"'
    
    # Find existing configuration file
    config_file = None
    for file in config_files:
        if os.path.exists(file):
            config_file = file
            break
    
    if not config_file:
        # If no configuration file found, create one
        config_file = config_files[0]
    
    print(f"  Updating {config_file}...")
    
    try:
        # Read existing content
        if os.path.exists(config_file):
            with open(config_file, 'r') as f:
                content = f.read()
        else:
            content = ""
        
        # Check if ANTHROPIC_BASE_URL is already set
        if 'ANTHROPIC_BASE_URL' in content:
            # Replace existing setting
            import re
            pattern = r'^export\s+ANTHROPIC_BASE_URL=.*$'
            new_content = re.sub(pattern, export_line, content, flags=re.MULTILINE)
            if new_content == content:
                # If no export statement matched, might be other format
                content += f"\n{export_line}\n"
            else:
                content = new_content
        else:
            # Add new setting
            if not content.endswith('\n'):
                content += '\n'
            content += f"\n# Claude Code Proxy Configuration\n{export_line}\n"
        
        # Write back to file
        with open(config_file, 'w') as f:
            f.write(content)
        
        print(f"  ✅ Environment variable added to {config_file}")
        print(f"  ⚠️ Please run the following command to activate the environment variable:")
        print(f"     source {config_file}")
        print(f"  Or reopen the terminal")
        
        # Set environment variable for current process
        os.environ['ANTHROPIC_BASE_URL'] = PROXY_BASE_URL
        print(f"  ✅ Environment variable set for current process")
        
        return True
        
    except Exception as e:
        print(f"  ❌ Failed to set environment variable: {e}")
        print(f"  Please manually add the following to your shell configuration file:")
        print(f"     {export_line}")
        return False

def main():
    """Main function"""
    print("=" * 60)
    print("Claude Code User Configuration Script")
    print("=" * 60)
    
    # Step 1: Check email parameter
    if len(sys.argv) != 2:
        print("❌ Error: Must provide email parameter")
        print("Usage: python3 setup_claude_user.py <email>")
        print("Example: python3 setup_claude_user.py user@example.com")
        sys.exit(1)
    
    email = sys.argv[1]
    
    if not validate_email(email):
        print(f"❌ Error: Invalid email format: {email}")
        sys.exit(1)
    
    print(f"📧 Configuring user: {email}")
    print()
    
    # Step 2: Check network connectivity
    if not check_connectivity():
        print("❌ Unable to connect to proxy server, please check VPN connection")
        sys.exit(1)
    
    print()
    
    # Step 3: Read Keychain password item
    print("Step 3: Checking credentials in Keychain...")
    credentials = get_keychain_credentials()
    
    if credentials:
        print("  ✅ Found Claude Code-credentials password item")
        
        # Check if it contains OAuth token
        access_token = None
        if isinstance(credentials, dict):
            # Check different possible structures
            if 'claudeAiOauth' in credentials:
                access_token = credentials['claudeAiOauth'].get('accessToken', '')
            elif 'accessToken' in credentials:
                access_token = credentials.get('accessToken', '')
            elif 'access_token' in credentials:
                access_token = credentials.get('access_token', '')
        
        if access_token and access_token.startswith('sk-ant'):
            print(f"  🔑 OAuth access token detected (sk-ant...)")
            print(f"  Token prefix: {access_token[:20]}...")
            
            # Add to OAuth token service
            print("  Adding OAuth token to service...")
            add_oauth_token(email, credentials)
        else:
            print("  ℹ️ Current password item content:")
            if isinstance(credentials, dict):
                print(f"     {json.dumps(credentials, indent=4)[:200]}...")
            else:
                print(f"     {str(credentials)[:200]}...")
    else:
        print("  ⚠️ Claude Code-credentials password item not found")
    
    print()
    
    # Step 4: Check and manage user list
    print("Step 4: Managing user list...")
    users = get_user_list()
    
    user_token = None
    user_exists = False
    
    # Check if email already exists
    for user in users:
        if user.get('name') == email:
            user_token = user.get('token')
            user_exists = True
            break
    
    if user_exists:
        print(f"  ✅ User already exists: {email}")
        print(f"  Token: {user_token[:20]}...")
    else:
        print(f"  ℹ️ User does not exist, creating new user...")
        user_token = generate_user_token()
        print(f"  Generated token: {user_token[:20]}...")
        
        # Add user
        user_result = add_user(email, user_token)
        if not user_result:
            print("  ⚠️ Failed to add user, but continuing with Keychain configuration...")
    
    print()
    
    # Step 5: Write to Keychain
    print("Step 5: Updating Keychain credentials...")
    
    if user_token:
        new_credentials = create_keychain_credentials(user_token)
        print("  Generating new credential structure:")
        print(f"     {json.dumps(json.loads(new_credentials), indent=4)[:300]}...")
        
        if write_to_keychain(new_credentials):
            print("  ✅ Successfully wrote to Keychain")
        else:
            print("  ❌ Failed to write to Keychain")
            sys.exit(1)
    else:
        print("  ❌ Unable to get user token")
        sys.exit(1)
    
    print()
    
    # Step 6: Check and set environment variable
    check_and_set_env_variable()
    
    print()
    print("=" * 60)
    print("✅ Configuration complete!")
    print(f"Claude Code environment has been set up for user {email}")
    print("=" * 60)

if __name__ == "__main__":
    main()